import pandas as pd
from flask import Blueprint, request, flash, redirect, url_for
from werkzeug.utils import secure_filename
from models import db, Material, Process, Product
import os

bp = Blueprint('excel_import', __name__)

ALLOWED_EXTENSIONS = {'xlsx', 'xls'}
UPLOAD_FOLDER = 'uploads'

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def ensure_upload_folder():
    if not os.path.exists(UPLOAD_FOLDER):
        os.makedirs(UPLOAD_FOLDER)

@bp.route('/import/materials', methods=['POST'])
def import_materials():
    if 'file' not in request.files:
        flash('没有选择文件')
        return redirect(url_for('materials'))
    
    file = request.files['file']
    if file.filename == '':
        flash('没有选择文件')
        return redirect(url_for('materials'))
    
    if file and allowed_file(file.filename):
        ensure_upload_folder()
        filename = secure_filename(file.filename)
        filepath = os.path.join(UPLOAD_FOLDER, filename)
        file.save(filepath)
        
        try:
            df = pd.read_excel(filepath)
            required_columns = ['材料名称', '规格', '品牌', '型号', '单位', '单价']
            if not all(col in df.columns for col in required_columns):
                flash('Excel文件格式不正确，请确保包含所有必需列')
                return redirect(url_for('materials'))
            
            for _, row in df.iterrows():
                material = Material(
                    name=row['材料名称'],
                    specification=row['规格'],
                    brand=row['品牌'],
                    model=row['型号'],
                    unit=row['单位'],
                    price=float(row['单价'])
                )
                db.session.add(material)
            
            db.session.commit()
            flash('材料数据导入成功')
            
        except Exception as e:
            db.session.rollback()
            flash(f'导入失败：{str(e)}')
        finally:
            os.remove(filepath)
        
        return redirect(url_for('materials'))
    
    flash('不支持的文件类型')
    return redirect(url_for('materials'))

@bp.route('/import/processes', methods=['POST'])
def import_processes():
    if 'file' not in request.files:
        flash('没有选择文件')
        return redirect(url_for('processes'))
    
    file = request.files['file']
    if file.filename == '':
        flash('没有选择文件')
        return redirect(url_for('processes'))
    
    if file and allowed_file(file.filename):
        ensure_upload_folder()
        filename = secure_filename(file.filename)
        filepath = os.path.join(UPLOAD_FOLDER, filename)
        file.save(filepath)
        
        try:
            df = pd.read_excel(filepath)
            required_columns = ['工序名称', '计价单位', '单价', '工序描述']
            if not all(col in df.columns for col in required_columns):
                flash('Excel文件格式不正确，请确保包含所有必需列')
                return redirect(url_for('processes'))
            
            for _, row in df.iterrows():
                process = Process(
                    name=row['工序名称'],
                    unit=row['计价单位'],
                    price=float(row['单价']),
                    description=row['工序描述']
                )
                db.session.add(process)
            
            db.session.commit()
            flash('工序数据导入成功')
            
        except Exception as e:
            db.session.rollback()
            flash(f'导入失败：{str(e)}')
        finally:
            os.remove(filepath)
        
        return redirect(url_for('processes'))
    
    flash('不支持的文件类型')
    return redirect(url_for('processes'))